import sqlparse
from sqlparse.tokens import Wildcard
from sqlparse.sql import Identifier, IdentifierList

class StarReplace:

    def starRelace(self,parsed,parsed_column_list):
        parsed_type = False #Identifier标志符
        for item in parsed:
            if 'WHERE' in item.value.upper():
                if 'SELECT' not in item.value.upper().split('WHERE')[0]:
                    continue

            if isinstance(parsed,Identifier):
                parsed_type = True
            if item.ttype == Wildcard and parsed_type == False:
                #处理 select *
                if parsed.value != '(*)':
                    #查看出现*号上一层是否是(*)
                    index = parsed.token_index(item)
                    parsed.parse_pop(index,parsed)
                    parsed.insert_before(index,parsed_column_list.pop(0))
            elif item.ttype == Wildcard and parsed_type == True:
                #处理 select t.*
                if parsed.value != '(*)':
                    parsed.parse_clear(parsed)
                    parsed.insert_before(0,parsed_column_list.pop(0))
            elif isinstance(item, Identifier) or isinstance(item, IdentifierList) or item.is_group:
                #其他情况需要递归下一层继续寻找*
                index = parsed.token_index(item)
                parsed.parse_pop(index, parsed)
                tmp_parsed = self.starRelace(item,parsed_column_list)
                parsed.insert_before(index,tmp_parsed)
        return parsed

    def handingColumn(self,sql,column_list):
        parsed_column_list = []  # 将传入的sql字段解析成parsed类型
        for column in column_list:
            if column[0] != '':
                parsed_column_list.append(sqlparse.parse(','.join(column))[0])
        parsed = sqlparse.parse(sql)[0]  # 解析完整的sql语句
        return self.starRelace(parsed, parsed_column_list)
